import os
from PIL import Image  # 图
import torch
from torch.utils.data import DataLoader  # 从 PIL 库中导入 Image 模块，用于图像处理。
from torchvision import models  # 导如一些已经训练好的模型
from model import load_model, feature_extract, transform
from dataset import MyDataset
import numpy as np
import faiss
from config import *




def img2feat(pic_file):  # 文件路径
    feat = []  # 建立空特征
    with open(pic_file, 'r', encoding='utf-8') as f:  # 打开文件夹，使用utf-8的方式
        lines = f.readlines()  # 读取文件的所有行，并将其存储在lines中
        # 从lines列表的第一行中提取特征向量值，并使用split()方法将其拆分为一个由浮点数字符串组成的列表
        # 然后，使用列表推导式将这些字符串转换为浮点数，并将其存储在feat列表中
        feat = [float(f) for f in lines[0].split()]
    return feat  # 返回获得的特征


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    # 载入模型
    model = load_model()


    # 图片库
    # img_folder = '../VOC2012_small'

    data_root = os.path.abspath(os.path.join(os.getcwd(), ".."))  # 拿到 souTU
    # img_folder = os.path.join(data_root, "VOC2012_small")  # 拿到 souTu/VOC2012_small
    img_folder = os.path.join(data_root, img_folder_name)  # 拿到 souTu/VOC2012_small

    assert os.path.exists(img_folder), "{} path does not exist.".format(img_folder)  # 断言警告
    # 创建一个 MyDataset 类的实例 val_dataset，将数据集文件夹的路径和 transform 对象传递给它。这样就得到了一个包含处理后的图像和相应索引值的数据集。
    val_dataset = MyDataset(img_folder, transform=transform)


    batch_size = 32  # 设置批处理的图像数量。
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print('Using {} dataloader workers every process'.format(nw))
    # 使用 DataLoader 类创建数据迭代器 val_dataloader
    val_dataloader = DataLoader(val_dataset,  # 需要加载的数据集
                                batch_size=batch_size,  # 批处理的大小
                                shuffle=False,  # 不需要对数据集进行随机打乱
                                num_workers=0)  # 使用 1 个进程加载数据集

    print("开始提取特征...")
    for idx, batch in enumerate(val_dataloader):  # idx 是指当前是多少个 batch
        img = batch['img']
        index = batch['index']

        # 解释：这里的 batch 是一个字典有两个键
        # print(batch.keys())  # dict_keys(['index', 'img'])
        # print(img.shape)     # torch.Size([64, 3, 224, 224])
        # print(index)         # tensor([0, 1, 2,...62, 63])

        img = img.to(device)  # 将图像数据image转移到指定设备上

        feature = feature_extract(model, img)  # 使用预训练模型model从图像数据img中提取图像特征feature。
        feature = feature.data.cpu().numpy()  # 将GPU中的张量数据转移到CPU中，并转换为Numpy数组类型。

        imgs_path = [os.path.join(img_folder, val_dataset.data[i] + '.txt') for i in index]
        # 使用图像索引index获取图像文件路径imgs_path，其中val_dataset.data[i]表示第i个图像的文件名（不包含扩展名）

        assert len(feature) == len(imgs_path)  # 断言特征数据feature的长度与图像文件路径imgs_path的长度相等。

        for i in range(len(imgs_path)):
            feature_list = [str(f) for f in feature[i]]  # 将图像特征feature转换为字符串列表feature_list，以便写入文件
            img_path = imgs_path[i]  # 获取当前图像的文件路径img_path

            with open(img_path, 'w', encoding='utf-8') as f:
                f.write(" ".join(feature_list))  # 将特征数据以空格分隔符连接为字符串，并写入文件

        print('*' * 60)  # 打印分割线。
        print(idx * batch_size)  # 打印当前批次数据的起始索引（第一个数据的索引）

    print("特征提取完毕")
    # =========== 至此我们得到了图片的特征并写入了文件 ============

    ids = []
    data = []

    img_path = os.path.join(img_folder, img_txt_name)  # 打开这个文件，这个文件都是文件名字的存储
    with open(img_path, 'r', encoding='utf-8') as f:  # 打开这个txt
        for line in f.readlines():  # 相当于遍历文本文件的每一行
            img_name = line.strip()  # 这一行代码的作用是读取文件 f 中的一行文本，然后去除字符串两端的空格，赋值给变量 img_name。该代码是读取图片文件名的操作。
            img_id = img_name.split('.')[0]  # 将图片名称（例如'cat.jpg'）按照"."进行分割，获取列表['cat', 'jpg']，然后取第一个元素'cat'作为图片的id。
            pic_txt_file = os.path.join(img_folder,
                                        "{}.txt".format(img_name))  # 这行代码是将当前图片的文件名和文件夹路径连接起来，构造出该图片所对应的标注文本文件的路径。

            if not os.path.exists(pic_txt_file):  # 如果图片的对应文本文件不存在，跳过当前循环，不执行后续代码。
                continue

            feat = img2feat(pic_txt_file)  # 得到特征
            ids.append(int(img_id))  # 得到图片id
            data.append(np.array(feat))  # 得到图片特征

    # 构建数据<id,data>
    ids = np.array(ids)  # 弄成数组，int
    data = np.array(data).astype('float32')  # 数组，float
    d = 512  # feature 特征长度（模型的结果）
    # print(" 特征向量记录数： ", data.shape)  # (1000, 512)：含有图片数据1000，其中每张图片记录了512个特征
    # print(" 特征向量ID的记录数：", ids.shape)  # (1000,)

    print("正在构建index_file.index...")
    # 创建图片特征索引－方案2(  资源有限，效果更好 )
    ###IDMap 支持add_with_ids
    ###如果很在意，使用”PCARx,...,SQ8“ 如果保存全部原始数据的开销太大，可以用这个索引方式。包含三个部分，
    # 1.降维
    # 2.聚类
    # 3.scalar 量化，每个向量编码为8bit 不支持GPU
    index = faiss.index_factory(d, "IDMap,PCAR16,IVF50,SQ8")  # 初始化，d表示特征向量的维度，这里为512，每个图像有512个特征数
    """
    IDMap：将向量 id 映射到其在列表中的位置，实现索引中向量 id 的快速检索。
    PCAR16：对输入的特征向量进行 PCA 降维，将其转换为长度为 16 的向量。这样可以减小内存占用和搜索时间，同时保留大部分特征向量的信息。
    IVF50：将向量空间划分为 50 个子空间（也称为“细分”），在每个子空间上建立一个倒排表（即 IVF），以存储在该子空间中的向量。这种方法可用于加速向量检索，特别是在高维空间中。
    SQ8：使用乘法量化器（SQ）将特征向量的每个分量量化为 8 个值中的一个。这可以大大减小内存占用，并加快相似度搜索。
    """
    index.train(data)  # 对数据集进行训练，以构建索引。
    index.add_with_ids(data, ids)  # 将数据集和ID添加到索引中。

    # 索引文件保存磁盘
    # 其中，faiss是用于高效相似度搜索和聚类的库，write_index是该库中用于将索引对象写入磁盘文件的方法。具体地，这里将索引对象index写入了名为index_file.index的文件
    faiss.write_index(index, index_file_name)
    # 这行代码的作用是从磁盘中读取索引文件index_file.index并创建一个faiss索引对象。该对象可用于对数据集进行快速搜索
    index = faiss.read_index(index_file_name)
    print(f"查看索引库大小:{index.ntotal}")  # 查看索引库大小


if __name__ == '__main__':
    main()
